

import akshare as ak
import pandas as pd
import numpy as np
from datetime import datetime, timedelta
import os


import logging



from logging.handlers import RotatingFileHandler



# 配置日志记录
logging.basicConfig(
    level=logging.INFO,  # 设置日志级别（DEBUG, INFO, WARNING, ERROR, CRITICAL）
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',  # 日志格式
    filename='choose_stock.log',  # 日志文件名
    filemode='a'  # 写入模式：'a'为追加，'w'为覆盖
)



# 创建文件处理器，指定编码为UTF-8
handler = logging.FileHandler('choose_stock.log', encoding='utf-8')


#832522
# 创建logger实例（可选，如果直接使用根logger可省略）
logger = logging.getLogger(__name__)

# 示例日志记录



# 获取当前工作目录（可能受运行方式影响）
current_working_dir = os.getcwd()
data_dir = current_working_dir + "/data/"

def calculate_technical_indicators(df):
    """
    计算技术指标
    """
    df = df.copy()
    
    # 5日收益率
    df['ret_5'] = df['收盘'].pct_change(5)
    
    # 5日成交量均线
    df['vol_ma5'] = df['成交量'].rolling(window=5).mean()
    
    # 20日均线
    df['ma20'] = df['收盘'].rolling(window=20).mean()
    
    # 成交量变化率
    df['vol_change'] = df['成交量'].pct_change()
    
    # 今日成交量与5日均量比值
    df['vol_ratio'] = df['成交量'] / df['vol_ma5']
    
    return df

def screen_stocks(date, stock_pool, lookback_days=30):
    """
    选股函数 - 在指定日期执行选股
    """
    end_date = date
    start_date = (datetime.strptime(date, '%Y-%m-%d') - timedelta(days=lookback_days+20)).strftime('%Y-%m-%d')
    
    selected_stocks = []
    
    for stock_code in stock_pool:
        try:
            #df = get_stock_data(stock_code, start_date, end_date)
            df_csv = pd.read_csv(os.path.join(data_dir, "stock", "{}.csv".format(stock_code)), encoding="utf_8_sig")
            df = df_csv.loc[df_csv['日期'] >= start_date]
            df = df.loc[df['日期'] <= end_date] 
            if df is None or len(df) < 30:  # 至少需要30天数据
                continue
                
            df = calculate_technical_indicators(df)
            
            # 获取最新一天的数据
            latest_data = df.iloc[-1]
            prev_data = df.iloc[-2]
            
            # 选股条件
            condition1 = latest_data['ret_5'] < -0.08  # 5日跌幅超过8%
            condition2 = latest_data['vol_ratio'] > 0.8  # 今日成交量不低于5日均量的80%
            condition3 = latest_data['vol_change'] > 0.3  # 成交量较前日放大30%
            condition4 = abs(latest_data['收盘'] - latest_data['ma20']) / latest_data['ma20'] < 0.02  # 接近20日均线
            logger.info(f"股票 {stock_code} - {latest_data['收盘']} - {latest_data['ret_5']} - {latest_data['vol_ratio']} - {latest_data['vol_change']} - {abs(latest_data['收盘'] - latest_data['ma20']) / latest_data['ma20']}")
            
            if condition1 and condition2 and condition3 and condition4:
                selected_stocks.append({
                    'code': stock_code,
                    'name': stock_code,  # 可额外获取股票名称
                    'close': latest_data['收盘'],
                    'ret_5': latest_data['ret_5'],
                    'signal_date': date
                })
                logger.info(f"选中池添加股票 {stock_code} - {latest_data['收盘']} - {latest_data['ret_5']} - {latest_data['vol_ratio']} - {latest_data['vol_change']} - {abs(latest_data['收盘'] - latest_data['ma20']) / latest_data['ma20']}")
                
        except Exception as e:

            logger.error(f"处理股票 {stock_code} 时出错: {e}")
            continue
    
    return selected_stocks

# 使用示例
if __name__ == "__main__":
    # 获取股票列表（示例，实际使用时需要完整的股票列表）
    stock_list  = pd.read_csv(os.path.join(data_dir, "stock_all_info.csv"), encoding="utf_8_sig")
    stock_all_info_wsfbs_pd = stock_list["代码"].str[2:] # 替换为完整的A股代码列表
    
    # 选股日期
    target_date = '2025-01-01'
    
    # 执行选股
    selected = screen_stocks(target_date, stock_all_info_wsfbs_pd)
    print(f"选中的股票: {selected}")
